from datetime import datetime
import pandas as pd
import time
from trade.data.table_edit import (
    add_factor_to_table,
    delete_factor_from_table,
    get_factor_from_table,
)
from trade.data.table_orm import session


def example_get_factor_from_table():
    """
    示例：调用 get_factor_from_table 从数据库中获取因子数据
    """
    print("Running example for `get_factor_from_table`...")

    # 定义查询参数
    timeframe = "5m"  # 1分钟的时间分辨率
    start = (
        datetime.fromisoformat("2023-10-01 08:00:00").timestamp() * 1000
    )  # 查询的开始时间
    end = (
        datetime.fromisoformat("2023-10-01 16:00:00").timestamp() * 1000
    )  # 查询的结束时间

    # 调用函数并获取结果
    result_df = get_factor_from_table(
        timeframe, start, end, symbols=["AAPL", "GOOG"], factors=["factorA", "factorB"]
    )

    # 打印查询结果
    print("Query result from database:")
    print(result_df)
    print("-" * 50)


def example_delete_factor_from_table():
    """
    示例：调用 delete_factor_from_table 删除特定条件下的数据
    """
    print("Running example for `delete_factor_from_table`...")

    # 定义删除条件
    timeframe = "5m"  # 5分钟的时间分辨率
    start = (
        datetime.fromisoformat("2023-10-01 08:00:00").timestamp() * 1000
    )  # 查询的开始时间
    end = (
        datetime.fromisoformat("2023-10-01 16:00:00").timestamp() * 1000
    )  # 查询的结束时间
    symbols = ["AAPL", "GOOG"]  # 指定删除的股票代码
    factors = ["factorA", "factorB"]  # 指定删除的因子名

    # 调用删除函数
    with session:
        delete_factor_from_table(timeframe, start, end, symbols, factors)
        session.commit()

    # 打印删除操作确认信息
    print(
        f"Deleted factors {factors} for symbols {symbols} "
        f"between {start} and {end}."
    )
    result_df = get_factor_from_table(
        timeframe, start, end, symbols=["AAPL", "GOOG"], factors=["factorA", "factorB"]
    )

    # 打印查询结果
    print("Query result from database:")
    print(result_df)

    print("-" * 50)


def example_add_factor_to_table():
    """
    示例：调用 add_factor_to_table 向数据库插入因子数据
    """
    print("Running example for `add_factor_to_table`...")

    # 构造插入的数据 DataFrame
    df_to_insert = pd.DataFrame(
        {
            "symbol": ["AAPL", "GOOG", "MSFT"],  # 股票代码
            "xdatetime": [
                datetime.fromisoformat("2023-10-01 08:00:00").timestamp()
                * 1000,  # 查询的开始时间
                datetime.fromisoformat("2023-10-01 09:00:00").timestamp()
                * 1000,  # 查询的结束时间
                time.time() * 1000,  # 查询的结束时间
            ],  # 时间戳
            "fname": ["factorA", "factorB", "factorC"],  # 因子名称
            "value": [777777.32, 2750.55, 295.45],  # 因子值
        }
    )
    df_to_insert["xdatetime"] = df_to_insert["xdatetime"].astype(int)
    print("Data to insert:")
    print(df_to_insert)

    timeframe = "5m"
    # print("最小最大时间", df_to_insert.xdatetime.min(), df_to_insert.xdatetime.max())
    with session:
        delete_factor_from_table(
            timeframe,
            start=df_to_insert.xdatetime.min(),
            end=df_to_insert.xdatetime.max(),
            symbols=df_to_insert.symbol.tolist(),
            factors=df_to_insert.fname.tolist(),
        )
        # 调用插入函数
        add_factor_to_table(df_to_insert, timeframe)
        session.commit()

    # 打印插入操作确认信息
    print(f"Added {len(df_to_insert)} rows to the {timeframe} table.")
    print("-" * 50)


def main():
    """
    主入口函数，用于运行上述所有的示例代码
    """
    print("Starting examples for `table_edit` functions...")
    print("=" * 50)

    # 按顺序运行每个示例函数
    example_get_factor_from_table()
    example_delete_factor_from_table()
    example_add_factor_to_table()

    print("All examples completed.")
    print("=" * 50)


if __name__ == "__main__":
    main()
